import os
import csv
import SimpleITK as sitk
import scipy.ndimage as ndi
import numpy as np
import nrrd

def calculate_HU(folder_path, image_path, csv_path, heart_niigz):
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        write = csv.writer(file)
        write.writerow(['file_name', 'total_HU', 'area'])
        for file in os.listdir(image_path):

            name = file.replace('nrrd', 'nii.gz')
            heart_path = os.path.join(heart_niigz, name)
            if not os.path.exists(heart_path):
                continue

            file_name1 = file.replace('.nrrd', 'heart_cac.nrrd')
            file_path = os.path.join(folder_path, file_name1)
            if not os.path.exists(file_path):
                continue

            try:
                image = sitk.ReadImage(file_path)
                image_array = sitk.GetArrayFromImage(image)
                old_image_path = os.path.join(image_path, file)
                print(old_image_path)
                image_old = sitk.ReadImage(old_image_path)
                image_array_old = sitk.GetArrayFromImage(image_old)
                total = 0.0
                area = 0.0
                for z in range(image_array.shape[0]):
                    # 计算面积
                    label_images, num_features = ndi.label(image_array[z, :, :])
                    slicer_total = 0
                    cac_area = 0
                    for label in range(1, num_features + 1):
                        mask = np.zeros_like(image_array[z, :, :])
                        mask[label_images == label] = 255
                        slice_data = mask
                        white_pixel_count = np.sum(slice_data == 255)
                        if white_pixel_count < 4:
                            continue
                        data, header = nrrd.read(old_image_path)
                        pixel_spacing = header.get('space directions', None)
                        pixel_area_slice = np.abs(pixel_spacing[0, 0] * pixel_spacing[1, 1] * pixel_spacing[2, 2]/3)
                        white_pixel_area = white_pixel_count * pixel_area_slice
                        # 计算weight
                        ct_data = image_array_old[z, :, :]
                        indices_255 = np.argwhere(slice_data == 255)
                        if len(indices_255) > 0:
                            hu_value = [ct_data[i, j] for i, j in indices_255]
                            max_hu_index = indices_255[np.argmax(hu_value)]
                        hu_value = np.abs(image_array_old[z, max_hu_index[0], max_hu_index[1]])
                        if 130 <= hu_value < 200:
                            weight = 1
                        elif 200 <= hu_value < 300:
                            weight = 2
                        elif 300 <= hu_value < 400:
                            weight = 3
                        elif hu_value >= 400:
                            weight = 4
                        slicer_total += weight * white_pixel_area
                        cac_area += white_pixel_area
                    total += slicer_total
                    area += cac_area
                print(total, area)
                write.writerow([file, total, area])
            except Exception as e:
                print(file)



def calculate_HU_ningshu(base_folder, csv_path):
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        write = csv.writer(file)
        write.writerow(['file_name', 'total_HU', 'area'])
        for folder in os.listdir(base_folder):
            folder_path = os.path.join(base_folder, folder)
            if len(os.listdir(folder_path)) != 3:
                continue

            name = 'heart.nrrd'
            heart_path = os.path.join(folder_path, name)
            if not os.path.exists(heart_path):
                continue

            file_name1 = 'label.nrrd'
            file_path = os.path.join(folder_path, file_name1)
            if not os.path.exists(file_path):
                continue

            try:
                image = sitk.ReadImage(file_path)
                image_array = sitk.GetArrayFromImage(image)
                old_image_path = os.path.join(folder_path, 'ori.nrrd')
                print(old_image_path)
                image_old = sitk.ReadImage(old_image_path)
                image_array_old = sitk.GetArrayFromImage(image_old)
                total = 0.0
                area = 0.0
                for z in range(image_array.shape[0]):
                    # 计算面积
                    label_images, num_features = ndi.label(image_array[z, :, :])
                    slicer_total = 0
                    cac_area = 0
                    for label in range(1, num_features + 1):
                        mask = np.zeros_like(image_array[z, :, :])
                        mask[label_images == label] = 255
                        slice_data = mask
                        white_pixel_count = np.sum(slice_data == 255)
                        if white_pixel_count < 4:
                            continue
                        data, header = nrrd.read(old_image_path)
                        pixel_spacing = header.get('space directions', None)
                        pixel_area_slice = np.abs(pixel_spacing[0, 0] * pixel_spacing[1, 1] * pixel_spacing[2, 2]/3)
                        white_pixel_area = white_pixel_count * pixel_area_slice
                        # 计算weight
                        ct_data = image_array_old[z, :, :]
                        indices_255 = np.argwhere(slice_data == 255)
                        if len(indices_255) > 0:
                            hu_value = [ct_data[i, j] for i, j in indices_255]
                            max_hu_index = indices_255[np.argmax(hu_value)]
                        hu_value = np.abs(image_array_old[z, max_hu_index[0], max_hu_index[1]])
                        if 130 <= hu_value < 200:
                            weight = 1
                        elif 200 <= hu_value < 300:
                            weight = 2
                        elif 300 <= hu_value < 400:
                            weight = 3
                        elif hu_value >= 400:
                            weight = 4
                        slicer_total += weight * white_pixel_area
                        cac_area += white_pixel_area
                    total += slicer_total
                    area += cac_area
                print(total, area)
                write.writerow([folder, total, area])
            except Exception as e:
                print(folder)


base_folder = '/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/test-ningshucac'
csv_path = '/media/imed/d9f520b2-6e7a-4842-af77-fba3ad2d8d27/test-ningshucac.csv'

calculate_HU_ningshu(base_folder, csv_path)

